import logging  
from sklearn.datasets import load_digits  
from sklearn.model_selection import train_test_split  
from tqdm import tqdm  
from pinecone import Pinecone, ServerlessSpec  
from collections import Counter 

# 配置日志  
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')  
  
# Pinecone 初始化  
pc = Pinecone(api_key="2183433c-f83f-4c83-9e03-6a4ae634ecfa")  
index_name = "mnist-index"  

# 获取现有索引列表
existing_indexes = pc.list_indexes()

# 检查索引是否存在，如果存在就删除
if any(index['name'] == index_name for index in existing_indexes):
    logging.info(f"索引 '{index_name}' 已存在，正在删除...")
    pc.delete_index(index_name)
    logging.info(f"索引 '{index_name}' 已成功删除。")
else:
    logging.infot(f"索引 '{index_name}' 不存在，将创建新索引。")

# 创建新索引
logging.info(f"正在创建新索引 '{index_name}'...")
pc.create_index(
    name=index_name,
    dimension=64,  # MNIST 每个图像展平后是一个 64 维向量
    metric="euclidean",  # 使用欧氏距离
    spec=ServerlessSpec(
        cloud="aws",
        region="us-east-1"
    )
)
logging.info(f"索引 '{index_name}' 创建成功。")

# 连接到索引
index = pc.Index(index_name)
logging.info(f"已成功连接到索引 '{index_name}'。")

# 加载MNIST数据集  
digits = load_digits(n_class=10)  
X = digits.data  
y = digits.target  
  
# 分割数据集  
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)  
  
# 转换数据为Pinecone可接受的格式  
def prepare_data(X, y):  
    vectors = []  
    for i, x in enumerate(X):  
        vector_id = str(i)  
        vector_values = x.tolist()  
        metadata = {"label": int(y[i])}  
        vectors.append((vector_id, vector_values, metadata))  
    return vectors  
  
# 准备训练数据和测试数据  
train_vectors = prepare_data(X_train, y_train)  
test_vectors = prepare_data(X_test, y_test)  
  
  
# 上传数据到索引  
index = pc.Index(index_name)  
batch_size = 1000 
for i in tqdm(range(0, len(train_vectors), batch_size), desc="Uploading data"):  
    batch = train_vectors[i:i + batch_size]  
    index.upsert(batch)  
logging.info(f"成功创建索引，并上传了 {len(train_vectors)} 条数据。")  
  
def evaluate_accuracy(index_name, test_vectors, k=11):  
    correct_predictions = 0  
    total_predictions = len(test_vectors)  
    for vector_id, vector_values, metadata in tqdm(test_vectors, desc="评估准确率"):  
        results = pc.Index(index_name).query(  
            vector=vector_values,  
            top_k=k,  
            include_metadata=True  
        )  
        labels = [match['metadata']['label'] for match in results['matches']]  
        if labels:  
            predicted_label = Counter(labels).most_common(1)[0][0]  
            if predicted_label == metadata['label']:  
                correct_predictions += 1  
    accuracy = correct_predictions / total_predictions  
    return accuracy  
# 计算准确率  
accuracy = evaluate_accuracy(index_name, test_vectors, k=11)  
logging.info(f"当k=11时，使用Pinecone的准确率为: {accuracy:.2f}")